package main

// go 实现一致性hash算法
// 算法原理：就是将数据通过hash的方式分发到不同的数据节点，当有节点挂了的时候只会影响少部分的数据迁移。
//          如果不通过一致性hash,通过简单的hash会将所有数据都需要重新迁移
// 迁移率能保持在1/n  公式：abs(原节点数-现有节点数)/ max(原节点数,现有节点数)
import (
"fmt"
"github.com/spaolacci/murmur3"
"log"
"reflect"
"sort"
"strconv"
"sync"
"unsafe"
)

// 自定义一个类型 hash环上的位置 或者hash环上数字的大小
type uints []uint32
// 给自定义类型绑定方法
func (x uints) Len() int { return len(x) }
func (x uints) Less(i, j int) bool { return x[i] < x[j] }
func (x uints) Swap(i, j int) { x[i], x[j] = x[j], x[i] }

// 定义一致性哈希环的数据结构
type Consistent struct {
	circle       map[uint32]string // 索引和key的map
	members      map[string]bool   // 快速查找key的set
	sortedHashes uints             // 索引的有序数组
	replicas     int               // 虚拟节点数
	count        int64             // 节点总数 实际节点总数
	sync.RWMutex                   //  读写锁
}

// 编写 初始化函数 传进去的就是虚拟节点数
func NewConsistent(replicas int) *Consistent {
	c := new(Consistent) // 返回结构体的指针
	c.replicas = replicas
	c.circle = make(map[uint32]string)
	c.members = make(map[string]bool)
	return c
}

// Add方法添加元素 key为物理服务器节点
func (c *Consistent) Add(key string) {
	c.Lock()
	defer c.Unlock()
	c.add(key)
}

// 内部的add方法，遍历虚拟节点数添加key
func (c *Consistent) add(key string) {
	// 遍历虚拟节点的次数
	for i := 0; i < c.replicas; i++ {
		c.circle[c.hashKey(c.genKey(key, i))] = key // 每个key有对应500个在circle
	}
	c.members[key] = true
	c.updateSortedHashes()
	c.count++
}

// 将key加上虚拟节点的索引
func (c *Consistent) genKey(key string, idx int) string {
	return strconv.Itoa(idx) + key // 字符串的拼接
}

// 根据key生成hash值，底层使用murmur3算法，这个算法比crc32均匀性要好
func (c *Consistent) hashKey(key string) uint32 {
	return murmur3.Sum32(str2bytes(key))
}

// 可以通过unsafe 强制转换绕过复制提高性能
func str2bytes(s string) (b []byte) {
	sh := *(*reflect.StringHeader)(unsafe.Pointer(&s))
	bh := (*reflect.SliceHeader)(unsafe.Pointer(&b))
	bh.Cap = sh.Len
	bh.Len = sh.Len
	bh.Data = sh.Data
	return b
}

// 对hash值切片进行排序
func (c *Consistent) updateSortedHashes() {
	// 取空切片
	hashes := c.sortedHashes[:0]
	for k := range c.circle {
		hashes = append(hashes, k) // 将所有的key放到hashes里
	}
	sort.Sort(hashes) // 一排序相当于所有的节点都打散了
	c.sortedHashes = hashes
}

// 根据一个target获取node
func (c *Consistent) Get(name string) string {
	c.RLock()
	defer c.RUnlock()
	key := c.hashKey(name)  // 相同name的hash获取到的key是一样的
	i := c.search(key)
	return c.circle[c.sortedHashes[i]] // 获取到对应的节点
}

// 索引切片的搜索方法，找到比key大的第一个值
func (c *Consistent) search(key uint32) (i int) {
	f := func(x int) bool {
		return c.sortedHashes[x] > key
	}
	i = sort.Search(len(c.sortedHashes), f)
	if i >= len(c.sortedHashes) {
		i = 0  // 防止下标越界
	}
	return
}

// 从hash环中删除一个节点
func (c *Consistent) Remove(elt string) {
	c.Lock()
	defer c.Unlock()
	c.remove(elt)
}

// 遍历虚拟节点数，获取key的哈希值，然后删除即可
func (c *Consistent) remove(key string) {
	for i := 0; i < c.replicas; i++ {
		// 根据key删除map的数据
		delete(c.circle, c.hashKey(c.genKey(key, i)))
	}
	delete(c.members, key)
	c.updateSortedHashes()
	c.count--
}

func main() {
	// 设置虚拟节点数500 先初始化一致性hash环
	r := NewConsistent(500)

	// 初始化的时候5个真实节点 key就是具体的物理节点
	nodes := []string{
		"1.1.1.1",
		"2.2.2.2",
		"3.3.3.3",
		"4.4.4.4",
		"5.5.5.5",
	}
	for _, n := range nodes {
		r.Add(n) // 5个节点入环 每一个物理节点key已经按照虚拟节点打散了
	}

	// 准备100个key
	keys := []string{}
	for i := 0; i < 1000; i++ {
		oneKey := fmt.Sprintf("%c_%d", i, i)
		keys = append(keys, oneKey)
	}
	// 100key获取哈希环中的节点
	keyNodeMap := make(map[string]string)
	for _, k := range keys {
		node := r.Get(k) // 获取该key对应的node
		keyNodeMap[k] = node
	}

	r.remove(nodes[1]) //移调一个node
	log.Printf("remove_node:%v", nodes[1])
	removeNum := 0
	for _, k := range keys {
		node := r.Get(k) // 继续计算key对应的node
		preNode := keyNodeMap[k] //该key没移除一个node前对应的node
		keyNodeMap[k] = node
		if preNode != node {
			//log.Printf("[key.node_change][key:%v][pre_node:%v][node:%v]", k, preNode, node)
			removeNum++ // 计算需要移动的数据
		}

	}
	log.Printf("[remove:%d/%d]", removeNum, len(keys))

	// 同理 添加一个node影响数据的计算
	r.add("6.6.6.6")
	log.Printf("add_node:%v", "6.6.6.6")
	removeNum = 0
	for _, k := range keys {
		node := r.Get(k)
		preNode := keyNodeMap[k]
		keyNodeMap[k] = node
		if preNode != node {
			//log.Printf("[key.node_change][key:%v][pre_node:%v][node:%v]", k, preNode, node)
			removeNum++
		}

	}
	log.Printf("[remove:%d/%d]", removeNum, len(keys))

}
